import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, r2_score
import warnings

warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


class SoilMoisturePredictor:
    #土壤湿度预测模型（含真实与预测数据可视化）

    def __init__(self):
        self.best_model = None
        self.features = ['T', 'Po', 'P', 'Pa', 'U', 'Ff', 'RRR']  # 表2指定气象因子
        self.min_soil_moisture = 0.22  # 附录1：作物最低存活湿度
        self.soil_dry_weight = 1500  # 附录3：单位土壤干重(kg/m³)

    def load_and_preprocess_data(self):

        print("=== 加载数据 ===")
        weather_data, soil_data = self._generate_mock_data()


        merged_data = pd.merge(
            weather_data,
            soil_data[['DateTime', '5cm_SM']],
            on='DateTime',
            how='inner'
        ).dropna(subset=self.features + ['5cm_SM'])

        print(f"有效样本量：{len(merged_data)}，土壤湿度范围："
              f"[{merged_data['5cm_SM'].min():.3f}, {merged_data['5cm_SM'].max():.3f}]")
        return merged_data

    def _generate_mock_data(self):

        np.random.seed(42)
        n_samples = 8760  # 5-7月逐时数据
        dates = pd.date_range('2021-05-01', periods=n_samples, freq='H')


        weather_data = pd.DataFrame({
            'DateTime': dates,
            'T': np.clip(22 + 6 * np.sin(np.linspace(0, 2*np.pi, n_samples)), 15, 35),
            'Po': np.clip(730 + np.random.normal(0, 2, n_samples), 720, 740),
            'P': np.clip(750 + np.random.normal(0, 2, n_samples), 740, 760),
            'Pa': np.clip(np.random.normal(0.5, 0.3, n_samples), -1, 2),
            'U': np.clip(65 + np.random.normal(0, 12, n_samples), 40, 100),
            'Ff': np.clip(np.random.normal(2, 1, n_samples), 0, 8),
            'RRR': np.maximum(0, np.random.exponential(0.3, n_samples))
        })


        soil_moisture = (
            0.25 +  # 基准湿度
            0.0005 * weather_data['U'] +
            0.008 * weather_data['RRR'] +
            -0.0025 * (weather_data['T'] - 20) +
            -0.001 * weather_data['Ff']
        ).clip(self.min_soil_moisture, 0.38)

        soil_data = pd.DataFrame({'DateTime': dates, '5cm_SM': soil_moisture})
        return weather_data, soil_data

    def train_model(self, data):
        #训练模型并可视化真实值与预测值对比
        print("\n=== 训练预测模型 ===")
        X = data[self.features]
        y = data['5cm_SM']

        # 划分数据集
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=0.2, random_state=42
        )

        # 训练模型
        self.best_model = RandomForestRegressor(
            n_estimators=120, max_depth=7, min_samples_split=10, random_state=42
        )
        self.best_model.fit(X_train, y_train)

        # 模型评估与预测
        y_pred = self.best_model.predict(X_test)
        print(f"模型性能：MSE={mean_squared_error(y_test, y_pred):.6f}, "
              f"R²={r2_score(y_test, y_pred):.4f}")

        # 可视化真实值与预测值（随机抽取100个样本）
        sample_idx = np.random.choice(len(y_test), 100, replace=False)
        plt.figure(figsize=(12, 6))
        plt.plot(y_test.iloc[sample_idx].values, label='真实土壤湿度', marker='o', linestyle='')
        plt.plot(y_pred[sample_idx], label='预测土壤湿度', marker='x', linestyle='')
        plt.axhline(y=self.min_soil_moisture, color='r', linestyle='--',
                   label=f'最低湿度阈值({self.min_soil_moisture})')
        plt.xlabel('样本索引')
        plt.ylabel('5cm土壤湿度')
        plt.title('真实土壤湿度与模型预测值对比（随机样本）')
        plt.legend()
        plt.tight_layout()
        plt.show()

        # 可视化预测误差分布
        plt.figure(figsize=(10, 6))
        sns.histplot(y_test - y_pred, kde=True)
        plt.axvline(x=0, color='r', linestyle='--')
        plt.xlabel('预测误差（真实值-预测值）')
        plt.title('土壤湿度预测误差分布')
        plt.tight_layout()
        plt.show()

        # 特征重要性分析
        feature_importance = pd.DataFrame({
            '气象因子': self.features,
            '重要性': self.best_model.feature_importances_
        }).sort_values('重要性', ascending=False)
        print("\n特征重要性：")
        print(feature_importance)

        plt.figure(figsize=(10, 6))
        sns.barplot(x='重要性', y='气象因子', data=feature_importance)
        plt.title('气象因子对土壤湿度的影响重要性')
        plt.tight_layout()
        plt.show()

        return self.best_model

    def predict_table2(self):
        #预测表2结果并输出
        if self.best_model is None:
            raise ValueError("请先训练模型")

        table2_input = pd.DataFrame({
            'T': [19.3, 20.0, 23.4, 28.0],
            'Po': [731.5, 732.0, 732.8, 733.5],
            'P': [751.7, 752.4, 753.2, 753.8],
            'Pa': [1.0, 0.5, 0.8, 0.7],
            'U': [99, 94, 80, 44],
            'Ff': [1.5, 6.0, 0, 0],
            'RRR': [0, 0, 0, 0]
        }, columns=self.features)

        predictions = self.best_model.predict(table2_input)
        predictions = np.clip(predictions, self.min_soil_moisture, 0.38)

        table2_result = pd.DataFrame({
            '时间': ['02:00', '05:00', '08:00', '11:00'],
            '5cm_SM（预测）': [round(p, 4) for p in predictions]
        })
        print("\n=== 表2土壤湿度预测结果 ===")
        print(table2_result)
        return table2_result


if __name__ == "__main__":
    predictor = SoilMoisturePredictor()
    data = predictor.load_and_preprocess_data()
    predictor.train_model(data)
    table2_result = predictor.predict_table2()